package com.firexxx.study.transformers;

import com.firexxx.study.visitor.TimeClassVisitor;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.nio.file.StandardOpenOption;
import java.security.ProtectionDomain;

/**
 * @description: TODO
 * @throws:
 * @author: zhuangzhang2
 * @date: 2020/07/02 12:27
 */
public class CustomVisitorTransformer implements ClassFileTransformer {
    private boolean detailSwitch = false;
    private String classPrefix = null;

    public CustomVisitorTransformer() {
    }

    public CustomVisitorTransformer(String classPrefix, boolean detailSwitch) {
        this.classPrefix = classPrefix;
        this.detailSwitch = detailSwitch;
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        String tmpClassName = className.replace("/", ".");
        if (classPrefix != null && !tmpClassName.startsWith(classPrefix)) {
            return null;
        }
        //读取类的字节码流
        ClassReader reader = new ClassReader(classfileBuffer);
        //创建操作字节流值对象，ClassWriter.COMPUTE_MAXS:表示自动计算栈大小
        ClassWriter writer = new ClassWriter(reader, ClassWriter.COMPUTE_MAXS);
        //接收一个ClassVisitor子类进行字节码修改
        reader.accept(new TimeClassVisitor(writer, className, detailSwitch), 8);
        //返回修改后的字节码流
        byte[] output = writer.toByteArray();

        if (classPrefix != null && tmpClassName.startsWith(classPrefix)) {
            try {
                Path path = Paths.get("output/classes/" + className + ".class");
                if (!Files.exists(path)) {
                    Files.createDirectories(path.getParent());
                }
                Files.write(path, output, StandardOpenOption.CREATE);
            } catch (Throwable t) {
                t.printStackTrace();
            }
        }

        return output;
    }

}